import numpy as np
import config
import sys


def process_results(gt, pr, classes):
    LEN = len(classes)
    matrix = np.zeros((LEN, LEN))
    for i, j in zip(gt, pr):
        matrix[i, j] += 1
    return matrix, classes


def list_err_info(id_str):
    gt_npy = './result/gt_' + id_str + '.npy'
    pr_npy = './result/pr_' + id_str + '.npy'
    sids_npy = './result/sids_' + id_str + '.npy'
    ts_npy = './result/ts_' + id_str + '.npy'

    # gt = np.load(config.gt_npy)
    # pr = np.load(config.pr_npy)
    gt = np.load(gt_npy)
    pr = np.load(pr_npy)
    sids = np.load(sids_npy)
    ts = np.load(ts_npy)
    err_idx = (1 - (gt == pr)).astype(bool)
    # print(err_idx)
    err_sids = sids[err_idx]
    err_ts = ts[err_idx]
    print("\nWrong predict info:")
    for sid, t in zip(err_sids, err_ts):
        print(sid, '\t', t)


def print_csv_with_id(id_str):
    if id_str is None:
        id_str = config.id_str
    print('  id_str', id_str)
    gt_npy = './result/gt_' + id_str + '.npy'
    pr_npy = './result/pr_' + id_str + '.npy'

    # gt = np.load(config.gt_npy)
    # pr = np.load(config.pr_npy)
    gt = np.load(gt_npy)
    pr = np.load(pr_npy)

    total_acc = np.sum(gt == pr) / len(gt)
    matrix, classes = process_results(gt, pr, config.emos)
    print()
    print('  a\\p', end='\t')
    for c in classes:
        print(c, end='\t')
    print()
    for i in range(len(classes)):
        print(' ', classes[i], end='\t')
        for ele in matrix[i]:
            print(ele, end='\t')
        print()
    print()

    sum_1 = np.sum(matrix, axis=1)
    matrix2 = matrix / sum_1.reshape((-1, 1))

    print('  a\\p', end='\t')
    for c in classes:
        print(' ', c, end='\t')
    print()
    for i in range(len(classes)):
        print(' ', classes[i], end='\t')
        for ele in matrix2[i]:
            print('%.4f' % ele, end='\t')
        print()
    print()

    avg = 0
    for i in range(len(classes)):
        avg += matrix2[i, i]
    print('  average(unweighted) accurate is %.4f' % (avg / len(classes)))
    print('  total(weighted) accurate is %.4f' % float(total_acc))
    print()
    if config.is_post_calc_len:
        list_err_info(id_str)


if __name__ == '__main__':
    if len(sys.argv) > 1:
        print_csv_with_id(sys.argv[1])
    else:
        print_csv_with_id(None)
